import os
import triton
from triton._filecheck import run_parser
from pathlib import Path


OUTPUT = os.getenv('OUTPUT', '.')


def env(fn, target):
    fname = fn._fn_name.replace('__main__.', '')

    print(f"Compile {fname} for {os.environ['GPU_NAME']} ({target})")

    os.environ['MLIR_ENABLE_DUMP'] = "1"
    os.environ['MLIR_DUMP_PATH'] = str(Path(OUTPUT + "/dump.mlir"))
    # os.environ['LLVM_IR_ENABLE_DUMP'] = "1"
    # os.environ['TRITON_ENABLE_LLVM_DEBUG'] = "1"
    # os.environ['TRITON_PRINT_AUTOTUNING'] = "1"
    os.environ['TRITON_ALWAYS_COMPILE'] = "1"
    os.environ['TRITON_KERNEL_DUMP'] = "1"
    os.environ['TRITON_DUMP_DIR'] = OUTPUT

    return fname


def save_ttir(fname, ttir):
    f = Path(OUTPUT + "/" + fname + ".ttir")
    f.write_text(str(ttir))
    return f


def compile_only(kernel, args, target):
    fn = kernel
    if isinstance(kernel, triton.runtime.autotuner.Autotuner):
        fn = kernel.fn

    fname = env(fn, target)

    t = run_parser(fn, args=args, target=target)
    f = save_ttir(fname, t)
    mod = triton.compile(src=str(f), target=target)
    return mod


def compile_ast_only(ast, target):
    fname = env(ast.fn, target)
    mod = triton.compile(src=ast, target=target)
    save_ttir(fname, mod.asm["ttir"])
    return mod
